// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/memory_pool.h>

#include <turbo/log/logging.h>

#include <benchmark/benchmark.h>

namespace nebula {

static constexpr int64_t kCacheLineSize = 64;

struct SystemAlloc {
  static turbo::Result<MemoryPool*> get_allocator() { return system_memory_pool(); }
};

static void TouchCacheLines(uint8_t* data, int64_t nbytes) {
  uint8_t total = 0;
  while (nbytes > 0) {
    total += *data;
    data += kCacheLineSize;
    nbytes -= kCacheLineSize;
  }
  benchmark::DoNotOptimize(total);
}

// Benchmark the cost of accessing always the same memory area.
// This gives us a lower bound of the potential difference between
// AllocateTouchDeallocate and AllocateDeallocate.
static void TouchArea(benchmark::State& state) {  // NOLINT non-const reference
  const int64_t nbytes = state.range(0);
  MemoryPool* pool = default_memory_pool();
  uint8_t* data;
  KCHECK_OK(pool->allocate(nbytes, &data));

  for (auto _ : state) {
    TouchCacheLines(data, nbytes);
  }

  pool->free(data, nbytes);
  state.SetItemsProcessed(state.iterations());
  state.SetBytesProcessed(state.iterations() * nbytes);
}

// Benchmark the raw cost of allocating memory.
// Note this is a best case situation: we always allocate and deallocate exactly
// the same size, without any other allocator traffic.  However, it can be
// representative of workloads where we routinely create and destroy
// temporary buffers for intermediate computation results.
template <typename Alloc>
static void AllocateDeallocate(benchmark::State& state) {  // NOLINT non-const reference
  const int64_t nbytes = state.range(0);
  MemoryPool* pool = *Alloc::get_allocator();

  for (auto _ : state) {
    uint8_t* data;
    KCHECK_OK(pool->allocate(nbytes, &data));
    pool->free(data, nbytes);
  }
  state.SetItemsProcessed(state.iterations());
  // SetBytesProcessed() would give nonsensical figures since the data is not
  // actually processed.
}

// Benchmark the cost of allocating memory plus accessing it.
template <typename Alloc>
static void AllocateTouchDeallocate(
    benchmark::State& state) {  // NOLINT non-const reference
  const int64_t nbytes = state.range(0);
  MemoryPool* pool = *Alloc::get_allocator();

  for (auto _ : state) {
    uint8_t* data;
    KCHECK_OK(pool->allocate(nbytes, &data));
    TouchCacheLines(data, nbytes);
    pool->free(data, nbytes);
  }
  state.SetItemsProcessed(state.iterations());
  state.SetBytesProcessed(state.iterations() * nbytes);
}

#define BENCHMARK_ALLOCATE_ARGS       \
  ->RangeMultiplier(16)               \
      ->Range(4096, 16 * 1024 * 1024) \
      ->ArgName("size")               \
      ->UseRealTime()                 \
      ->ThreadRange(1, 32)

#define BENCHMARK_ALLOCATE(benchmark_func, template_param) \
  BENCHMARK_TEMPLATE(benchmark_func, template_param) BENCHMARK_ALLOCATE_ARGS

BENCHMARK(TouchArea) BENCHMARK_ALLOCATE_ARGS;

BENCHMARK_ALLOCATE(AllocateDeallocate, SystemAlloc);
BENCHMARK_ALLOCATE(AllocateTouchDeallocate, SystemAlloc);



}  // namespace nebula
